import numpy as np
from ultralytics import YOLO
from dataclasses import dataclass
import typing
from torch import Tensor
import cv2


@dataclass
class Box:
    x: int
    y: int
    width: int
    height: int


@dataclass
class RoiBox:
    box: Box
    confidence: float
    class_id: int


class YoloDetector:
    def __init__(self, weight_path='/home/hello/下载/yolo11x-seg.pt'):
        # Load a model
        self.model = YOLO(weight_path)  # pretrained YOLO11n model

    def detect(self, image: np.ndarray) -> (typing.List[RoiBox], np.ndarray):
        def to_box(t: Tensor) -> Box:
            t = t.int().tolist()
            assert len(t) == 4
            return Box(*t)

        # Run batched inference on a list of images
        result = self.model.predict(image)[0]  # return a list of Results objects
        # 假设你已经有了 result 对象
        annotated_image = result.plot()  # 这会返回一个 numpy 数组
        cv2_image = cv2.cvtColor(annotated_image, cv2.COLOR_RGB2BGR)
        roi_list = []
        # Process results list
        for box, confidence, class_id in zip(result.boxes.xywh, result.boxes.conf, result.boxes.cls):
            box_ = to_box(box)
            roi_box = RoiBox(box_, float(confidence.cfloat()), int(class_id.int()))
            roi_list.append(roi_box)
        return roi_list, cv2_image, result.names
